package org.cybzacg.encryption.impl;

import org.cybzacg.encryption.core.AbstractEncryptor;
import org.cybzacg.encryption.context.EncryptionContext;
import org.cybzacg.encryption.model.EncryptionResult;
import org.cybzacg.encryption.model.ValidationResult;
import org.cybzacg.encryption.enums.EncryptionAlgorithm;
import org.cybzacg.encryption.enums.EncryptionType;
import org.cybzacg.encryption.enums.NationalAlgorithm;
import org.cybzacg.encryption.utils.ValidationUtils;
import org.cybzacg.encryption.utils.PerformanceUtils;

import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.*;
import java.util.Arrays;

/**
 * SM4加密器实现
 * 中国商用密码分组密码算法
 *
 * @author encryption team
 * @version 1.0
 */
public class SM4Encryptor extends AbstractEncryptor {

    private static final String ALGORITHM_NAME = "SM4";
    private static final String TRANSFORMATION = "SM4/CBC/PKCS7Padding";
    private static final int BLOCK_SIZE = 16; // SM4块大小128位=16字节
    private static final int KEY_SIZE = 16; // SM4密钥大小128位=16字节

    /**
     * 构造函数
     */
    public SM4Encryptor() {
        super(EncryptionAlgorithm.SM4);
    }

    @Override
    protected byte[] doEncrypt(byte[] data, EncryptionContext context) {
        try (PerformanceUtils.PerformanceContext performanceContext =
                PerformanceUtils.startMonitoring("SM4 encryption")) {

            // 验证输入参数
            ValidationResult validationResult = validateEncryptionInput(data, context);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 获取密钥
            SecretKey key = getSecretKey(context);
            if (key == null) {
                throw new IllegalArgumentException("SM4 key not available in context");
            }

            // 获取或生成IV
            byte[] iv = getIV(context);
            if (iv == null) {
                iv = generateIV();
            // 将IV保存到上下文中，以便解密使用
            context.setAttribute("iv", iv);
            }

            // 执行加密
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            IvParameterSpec ivSpec = new IvParameterSpec(iv);
            cipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);

            return cipher.doFinal(data);

        } catch (Exception e) {
            throw new RuntimeException("SM4 encryption failed: " + e.getMessage(), e);
        }
    }

    @Override
    protected byte[] doDecrypt(byte[] encryptedData, EncryptionContext context) {
        try (PerformanceUtils.PerformanceContext performanceContext =
                PerformanceUtils.startMonitoring("SM4 decryption")) {

            // 验证输入参数
            ValidationResult validationResult = validateDecryptionInput(encryptedData, context);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 获取密钥
            SecretKey key = getSecretKey(context);
            if (key == null) {
                throw new IllegalArgumentException("SM4 key not available in context");
            }

            // 获取IV
            byte[] iv = getIV(context);
            if (iv == null) {
                // 尝试从上下文中获取IV
                Object ivObj = context.getAttribute("iv");
                if (ivObj instanceof byte[]) {
                    iv = (byte[]) ivObj;
                } else {
                    throw new IllegalArgumentException("IV not available in context");
                }
            }

            // 执行解密
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            IvParameterSpec ivSpec = new IvParameterSpec(iv);
            cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);

            return cipher.doFinal(encryptedData);

        } catch (Exception e) {
            throw new RuntimeException("SM4 decryption failed: " + e.getMessage(), e);
        }
    }

    @Override
    protected void doInitialize(EncryptionContext context) {
        // SM4特定的初始化逻辑
        if (context != null) {
            // 可以在这里添加SM4特定的初始化参数
            context.setAttribute("algorithm", ALGORITHM_NAME);
            context.setAttribute("transformation", TRANSFORMATION);
            context.setAttribute("blockSize", BLOCK_SIZE);
            context.setAttribute("keySize", KEY_SIZE);
        }
    }

    @Override
    protected ValidationResult doValidate(EncryptionContext context) {
        if (context == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Context cannot be null")
                    .build();
        }

        // 验证密钥
        SecretKey key = getSecretKey(context);
        if (key == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 key not available in context")
                    .build();
        }

        return validateKey(key, context);
    }

    @Override
    protected boolean doSupports(EncryptionContext context) {
        if (context == null) {
            return false;
        }

        // 检查算法匹配
        if (getAlgorithm() != context.getAlgorithm()) {
            return false;
        }

        // 检查加密类型匹配
        if (getType() != context.getEncryptionType()) {
            return false;
        }

        return true;
    }

    /**
     * 验证密钥
     */
    public ValidationResult validateKey(SecretKey key, EncryptionContext context) {
        if (key == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 key cannot be null")
                    .build();
        }

        byte[] keyBytes = key.getEncoded();
        if (keyBytes == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 key encoding cannot be null")
                    .build();
        }

        if (keyBytes.length != KEY_SIZE) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM4 key must be exactly %d bytes (%d bits), got %d bytes",
                        KEY_SIZE, KEY_SIZE * 8, keyBytes.length))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM4 key is valid")
                .build();
    }


    /**
     * 验证加密输入参数
     */
    private ValidationResult validateEncryptionInput(byte[] data, EncryptionContext context) {
        // 验证数据
        if (data == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 encryption data cannot be null")
                    .build();
        }

        if (data.length == 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 encryption data cannot be empty")
                    .build();
        }

        // 验证数据长度
        if (data.length > getMaxDataSize()) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM4 encryption data size %d bytes exceeds maximum %d bytes",
                        data.length, getMaxDataSize()))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM4 encryption input is valid")
                .build();
    }

    /**
     * 验证解密输入参数
     */
    private ValidationResult validateDecryptionInput(byte[] encryptedData, EncryptionContext context) {
        // 验证加密数据
        if (encryptedData == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 encrypted data cannot be null")
                    .build();
        }

        if (encryptedData.length == 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM4 encrypted data cannot be empty")
                    .build();
        }

        // 验证加密数据长度
        if (encryptedData.length > getMaxDataSize()) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM4 encrypted data size %d bytes exceeds maximum %d bytes",
                        encryptedData.length, getMaxDataSize()))
                    .build();
        }

        // 验证加密数据长度是否为块大小的倍数
        if (encryptedData.length % BLOCK_SIZE != 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM4 encrypted data length %d bytes must be a multiple of block size %d bytes",
                        encryptedData.length, BLOCK_SIZE))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM4 decryption input is valid")
                .build();
    }

    /**
     * 获取密钥
     */
    private SecretKey getSecretKey(EncryptionContext context) {
        // 尝试从上下文中获取密钥
        Object keyObj = context.getAttribute("secretKey");
        if (keyObj instanceof SecretKey) {
            return (SecretKey) keyObj;
        }

        // 尝试从上下文中获取密钥字节数组
        Object keyBytesObj = context.getAttribute("keyBytes");
        if (keyBytesObj instanceof byte[]) {
            byte[] keyBytes = (byte[]) keyBytesObj;
            try {
                return new SecretKeySpec(keyBytes, ALGORITHM_NAME);
            } catch (Exception e) {
                return null;
            }
        }

        return null;
    }

    /**
     * 获取IV
     */
    private byte[] getIV(EncryptionContext context) {
        // 尝试从上下文中获取IV
        Object ivObj = context.getAttribute("iv");
        if (ivObj instanceof byte[]) {
            return (byte[]) ivObj;
        }

        return null;
    }

    /**
     * 生成IV
     */
    private byte[] generateIV() {
        byte[] iv = new byte[BLOCK_SIZE];
        new SecureRandom().nextBytes(iv);
        return iv;
    }

    /**
     * 获取最大数据大小
     */
    private int getMaxDataSize() {
        // SM4最大加密数据大小限制
        return 1024 * 1024; // 1MB
    }

    /**
     * 生成SM4密钥
     *
     * @return SM4密钥
     */
    public static SecretKey generateKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM_NAME);
            keyGenerator.init(KEY_SIZE * 8);
            return keyGenerator.generateKey();
        } catch (Exception e) {
            throw new RuntimeException("Failed to generate SM4 key: " + e.getMessage(), e);
        }
    }

    /**
     * 从字节数组创建SM4密钥
     *
     * @param keyBytes 密钥字节数组
     * @return SM4密钥
     */
    public static SecretKey createKey(byte[] keyBytes) {
        if (keyBytes == null) {
            throw new IllegalArgumentException("Key bytes cannot be null");
        }

        if (keyBytes.length != KEY_SIZE) {
            throw new IllegalArgumentException(
                String.format("SM4 key must be exactly %d bytes, got %d bytes",
                    KEY_SIZE, keyBytes.length)
            );
        }

        return new SecretKeySpec(keyBytes, ALGORITHM_NAME);
    }

    /**
     * 生成随机IV
     *
     * @return IV字节数组
     */
    public static byte[] generateRandomIV() {
        byte[] iv = new byte[BLOCK_SIZE];
        new SecureRandom().nextBytes(iv);
        return iv;
    }

    /**
     * 检查SM4算法是否可用
     *
     * @return 是否可用
     */
    public static boolean isAvailable() {
        try {
            Cipher.getInstance(TRANSFORMATION);
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 获取算法信息
     *
     * @return 算法信息
     */
    public static String getAlgorithmInfo() {
        return String.format(
            "SM4 Algorithm Information:\n" +
            "  Algorithm: %s\n" +
            "  Transformation: %s\n" +
            "  Block Size: %d bytes\n" +
            "  Key Size: %d bytes\n" +
            "  Available: %s",
            ALGORITHM_NAME, TRANSFORMATION, BLOCK_SIZE, KEY_SIZE, isAvailable()
        );
    }

    /**
     * SM4密钥生成器
     */
    public static class SM4KeyGenerator {

        /**
         * 生成随机SM4密钥
         *
         * @return SM4密钥
         */
        public static SecretKey generateRandomKey() {
            return generateKey();
        }

        /**
         * 从密码生成SM4密钥
         *
         * @param password 密码
         * @param salt 盐值
         * @param iterations 迭代次数
         * @return SM4密钥
         */
        public static SecretKey deriveKey(char[] password, byte[] salt, int iterations) {
            try {
                SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
                PBEKeySpec spec = new PBEKeySpec(
                    password, salt, iterations, KEY_SIZE * 8);
                SecretKey tmp = factory.generateSecret(spec);
                return new SecretKeySpec(tmp.getEncoded(), ALGORITHM_NAME);
            } catch (Exception e) {
                throw new RuntimeException("Failed to derive SM4 key from password: " + e.getMessage(), e);
            }
        }

        /**
         * 从种子生成SM4密钥
         *
         * @param seed 种子
         * @return SM4密钥
         */
        public static SecretKey generateKeyFromSeed(byte[] seed) {
            try {
                MessageDigest md = MessageDigest.getInstance("SHA-256");
                byte[] hash = md.digest(seed);
                return new SecretKeySpec(Arrays.copyOf(hash, KEY_SIZE), ALGORITHM_NAME);
            } catch (Exception e) {
                throw new RuntimeException("Failed to generate SM4 key from seed: " + e.getMessage(), e);
            }
        }
    }

    /**
     * SM4参数生成器
     */
    public static class SM4ParameterGenerator {

        /**
         * 生成随机参数
         *
         * @return 参数映射
         */
        public static java.util.Map<String, Object> generateRandomParameters() {
            java.util.Map<String, Object> params = new java.util.HashMap<>();
            params.put("iv", generateRandomIV());
            return params;
        }

        /**
         * 生成固定参数
         *
         * @param seed 种子
         * @return 参数映射
         */
        public static java.util.Map<String, Object> generateFixedParameters(byte[] seed) {
            java.util.Map<String, Object> params = new java.util.HashMap<>();
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("SHA-256");
                byte[] hash = md.digest(seed);
                byte[] iv = Arrays.copyOf(hash, BLOCK_SIZE);
                params.put("iv", iv);
            } catch (Exception e) {
                throw new RuntimeException("Failed to generate fixed parameters: " + e.getMessage(), e);
            }
            return params;
        }
    }
}
